package dd.lo.chapter11;

import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.photo.Photo;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 代码样例：Image repair by inpainting（可处理其他颜色的少量缺失部分）
 */
public class Example5 extends Application {

    public static void start(String[] args) {
        launch(args);
    }

    private enum Mode {
        ORIGIN,
        MASK,
        PREVIEW,
        INPAINT,
    }

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    private static final int COLOR_OFFSET = 15;

    private ImageView imageView;

    private Mat srcImg;

    private Mat img;

    private int[] point;

    private void displayImage(Mode mode) {
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() {
                if (null == srcImg) return null;
                Mat dst = new Mat();
                srcImg.copyTo(dst);
                ImageUtils.fitImgSize(dst, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                if (null == img) {
                    img = new Mat();
                    srcImg.copyTo(img);
                }
                if (null != point) {
                    if (Mode.MASK == mode) {
                        double[] pixel = dst.get(point[1], point[0]);
                        System.out.printf("Click point [%d, %d] pixel: [%f, %f, %f]\n", point[0], point[1], pixel[0], pixel[1], pixel[2]);
                        Mat mask = new Mat();
                        //Output is a single channel
                        Core.inRange(dst, new Scalar(pixel[0] - COLOR_OFFSET, pixel[1] - COLOR_OFFSET, pixel[2] - COLOR_OFFSET), new Scalar(pixel), mask);
                        //Review mask
                        mask.copyTo(dst);
                    } else if (Mode.INPAINT == mode) {
                        double[] pixel = dst.get(point[1], point[0]);
                        System.out.printf("Click point [%d, %d] pixel: [%f, %f, %f]\n", point[0], point[1], pixel[0], pixel[1], pixel[2]);
                        Mat mask = new Mat();
                        //Output is a single channel
                        Core.inRange(img, new Scalar(pixel[0] - COLOR_OFFSET, pixel[1] - COLOR_OFFSET, pixel[2] - COLOR_OFFSET), new Scalar(pixel), mask);
                        //Strengthen mask by threshold and dilate
                        Imgproc.threshold(mask, mask, 254, 255, Imgproc.THRESH_TOZERO);
                        Imgproc.dilate(mask, mask, new Mat(), new Point(-1, -1), 5);
                        //Apply inpaint function
                        Photo.inpaint(img, mask, img, 10, Photo.INPAINT_NS);
                        img.copyTo(dst);
                        ImageUtils.fitImgSize(dst, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                    } else if (Mode.PREVIEW == mode) {
                        img.copyTo(dst);
                        ImageUtils.fitImgSize(dst, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                    } else {
                        //Mode.ORIGIN == mode
                        if (null == img) {
                            img = new Mat();
                        }
                        srcImg.copyTo(img);
                    }
                }
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", dst, buffer);
                Platform.runLater(() -> imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray()))));
                return null;
            }
        };
        new Thread(task).start();
    }

    private void saveImage(File file) {
        if (null == img) {
            System.out.println("还没有处理结果图片");
            return;
        }
        if (!file.exists()) {
            try {
                boolean created = file.createNewFile();
                if (!created) throw new IOException("生成文件失败");
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        boolean written = Imgcodecs.imwrite(file.getAbsolutePath(), img);
        System.out.println(written ? "写出文件成功" : "写出文件失败");
    }

    @Override
    public void start(Stage stage) {
        Button imageChooser = new Button("选择图片");
        imageChooser.setOnMouseClicked(e -> {
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择图片");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            fileChooser.getExtensionFilters().addAll(
//                    new FileChooser.ExtensionFilter("All Images", "*.*"),
                    new FileChooser.ExtensionFilter("JPEG", "*.jpeg"),
                    new FileChooser.ExtensionFilter("WEBP", "*.webp"),
                    new FileChooser.ExtensionFilter("PNG", "*.png"),
                    new FileChooser.ExtensionFilter("JPG", "*.jpg"),
                    new FileChooser.ExtensionFilter("GIF", "*.gif"),
                    new FileChooser.ExtensionFilter("BMP", "*.bmp")
            );
            File imgFile = fileChooser.showOpenDialog(stage);
            if (null == imgFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //读取一张图片
            srcImg = Imgcodecs.imread(imgFile.getAbsolutePath(), Imgcodecs.IMREAD_UNCHANGED);
            if (srcImg.empty()) {
                throw new RuntimeException("图片读取失败");
            }
            displayImage(Mode.ORIGIN);
        });
        //Origin button
        Button originButton = new Button("原图");
        originButton.setOnMouseClicked(e -> {
            displayImage(Mode.ORIGIN);
        });
        //Unselect button
        Button unselectBtn = new Button("取消选中");
        unselectBtn.setOnMouseClicked(e -> {
            displayImage(Mode.PREVIEW);
        });
        //Inpaint button
        Button inpaintButton = new Button("修复");
        inpaintButton.setOnMouseClicked(e -> {
            displayImage(Mode.INPAINT);
        });
        //Save button
        Button saveImgBtn = new Button("保存图片");
        saveImgBtn.setOnMouseClicked(e -> {
            FileChooser fileChooser = new FileChooser();
            fileChooser.setTitle("请选择文件保存位置");
            fileChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            File imgFile = fileChooser.showSaveDialog(stage);
            if (null == imgFile) {
                System.out.println("用户取消选中文件");
                return;
            }
            //选中保存位置文件
            saveImage(imgFile);
        });
        //Exit button
        Button exitButton = new Button("关闭");
        exitButton.setOnMouseClicked(e -> {
            Platform.exit();
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(imageChooser, originButton, unselectBtn, inpaintButton, saveImgBtn, exitButton);
        BorderPane root = new BorderPane();
        imageView = new ImageView();
        imageView.setOnMouseClicked(e -> {
            //Inpaint by the color indicated by the mouse click
            point = new int[]{(int) e.getX(), (int) e.getY()};
            displayImage(Mode.MASK);
        });
        root.setCenter(imageView);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("Example5 - Image repair on color indicated small scratches");
        stage.setResizable(false);
        stage.setOnCloseRequest(e -> {
            System.out.println("Receive close request");
            e.consume();
        });
        stage.show();
    }
}
